package com.lap.dev.starter.core;

import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

@Slf4j
public class CustomLoadBalancer implements ReactorServiceInstanceLoadBalancer {

  public static final String SHARED = "shared";
  public static final String OWNER = "owner";

  private final String serviceId;
  private final ObjectProvider<ServiceInstanceListSupplier> objectProvider;
  private final Set<String> localIps;

  private final Random random = new Random();

  public CustomLoadBalancer(
      String serviceId,
      ObjectProvider<ServiceInstanceListSupplier> supplierProvider,
      Set<String> localIpSet) {
    this.serviceId = serviceId;
    this.objectProvider = supplierProvider;
    this.localIps = localIpSet;
  }

  @Override
  public Mono<Response<ServiceInstance>> choose(Request request) {
    ServiceInstanceListSupplier supplier = objectProvider.getIfAvailable();
    if (supplier == null) {
      log.warn("No instance supplier available for service>>: {}", serviceId);
      return Mono.just(new EmptyResponse());
    }

    return supplier.get().next().map(this::selectInstance);
  }

  private Response<ServiceInstance> selectInstance(List<ServiceInstance> instances) {
    if (instances.isEmpty()) {
      log.warn("No service instances available for service>>: {}", serviceId);
      return new EmptyResponse();
    }

    // 阶段1: 匹配本地IP实例
    List<ServiceInstance> localMatches =
        instances.stream()
            .filter(instance -> localIps.contains(instance.getHost().toLowerCase()))
            .collect(Collectors.toList());

    if (!localMatches.isEmpty()) {
      log.debug("Found {} local instances for service: {}", localMatches.size(), serviceId);
      return new DefaultResponse(selectRandom(localMatches));
    }

    // 阶段2: 匹配Owner实例
    List<ServiceInstance> ownerMatches =
        instances.stream()
            .filter(
                instance ->
                    localIps.contains(instance.getMetadata().getOrDefault(OWNER, "").toLowerCase()))
            .collect(Collectors.toList());

    if (!ownerMatches.isEmpty()) {
      log.debug("Found {} owner-matched instances", ownerMatches.size());
      return new DefaultResponse(selectRandom(ownerMatches));
    }

    // 阶段3: 匹配共享实例
    List<ServiceInstance> sharedMatches =
        instances.stream()
            .filter(instance -> SHARED.equalsIgnoreCase(instance.getMetadata().get(OWNER)))
            .collect(Collectors.toList());

    if (!sharedMatches.isEmpty()) {
      log.debug("Found {} shared instances", sharedMatches.size());
      return new DefaultResponse(selectRandom(sharedMatches));
    }

    // 兜底: 加权随机选择
    log.debug("Fallback to global random selection");
    return new DefaultResponse(selectRandom(instances));
  }

  private ServiceInstance selectRandom(List<ServiceInstance> instances) {
    return instances.get(random.nextInt(instances.size()));
  }
}
